
import numpy as np
import librosa
import os
from python_speech_features import logfbank
from patient_information import load_patient_data,get_Murmur_locations
from python_speech_features import fbank

def FrameTimeC(frameNum, frameLen, inc, fs):
    ll = np.array([i for i in range(frameNum)])
    return ((ll - 1) * inc + frameLen / 2) / fs
def FrequencyScale(nfilt,fs):
    high_freq_mel = (2595 * np.log10(1 + (fs / 2) / 700))  # 求最高hz频率对应的mel频率
# 我们要做40个滤波器组，为此需要42个点，这意味着在们需要low_freq_mel和high_freq_mel之间线性间隔40个点
    mel_points = np.linspace(0, high_freq_mel, nfilt + 2)  # 在mel频率上均分成42个点
    hz_points = (700 * (10 ** (mel_points / 2595) - 1))  # 将mel频率再转到hz频率
    return hz_points

def Log_Mel_Spec(data_directory ):
    logmelspec = list()
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            x, fs = librosa.load(os.path.join(data_directory, f), sr=4000)
            fbank_feat = logfbank(x,fs,winlen=0.025,winstep=0.0125,nfilt=32,nfft=512,lowfreq=0,highfreq=800)
            fbank_feat=fbank_feat.T
            logmelspec.append(fbank_feat)
        else:
            continue
    return np.array(logmelspec)

def get_label(data_directory ):
    label=list()
    location=list()
    id=list()
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            the_location = root.split("_")[1].strip()
            location.append(the_location)

            the_id = root.split("_")[0].strip()
            id.append(the_id)

            #将不存在杂音的听诊位置标为absent
            the_label='Absent'
            txt_data = load_patient_data(os.path.join(data_directory,the_id+'.txt'))
            murmur_locations = (get_Murmur_locations(txt_data)).split("+")
            for i in range(len(murmur_locations)):
                if the_location == murmur_locations[i]:
                    the_label = root.split("_")[2].strip()

            if the_label == 'Absent':
                grade = 0
            elif the_label == 'Soft':
                grade = 1
            elif the_label == 'Loud':
                grade = 2
            label.append(grade)

    return np.array(label),np.array(location),np.array(id)

def get_index(data_directory ):
    index = list()
    i=0
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            index.append(i)
            i=i+1
        else:
            continue
    return np.array(index)



if __name__ == "__main__":

    vali_data_directory="data/stratified_data_new/vali_data"
    test_data_directory = "data/stratified_data_new/test_data"
    train_data_directory = "data/stratified_data_new/train_data"
    out_directory="data/logmel"
    label_directory="data/label"
    # if not os.path.exists(out_directory):
    #     os.makedirs(out_directory)
    # if not os.path.exists(label_directory):
    #     os.makedirs(label_directory)
    # #提取logmel特征 并保存
    # train_feature = Log_Mel_Spec(train_data_directory)
    # vali_feature = Log_Mel_Spec(vali_data_directory)
    # test_feature = Log_Mel_Spec(test_data_directory)
    # np.save(out_directory+r'/train_feature.npy',train_feature)
    # np.save(out_directory + r'/vali_feature.npy', vali_feature)
    # np.save(out_directory + r'/test_feature.npy', test_feature)
    # #提取label 并保存
    train_label , train_location,train_id= get_label(train_data_directory)
    vali_label ,vali_location,vali_id= get_label(vali_data_directory)
    test_label , test_location,test_id = get_label(test_data_directory)
    np.save(label_directory + r'/train_label.npy',train_label)
    np.save(label_directory + r'/vali_label.npy', vali_label)
    np.save(label_directory + r'/test_label.npy', test_label)
    # #保存每个wav的听诊区
    # np.save(label_directory + r'/train_location.npy',train_location)
    # np.save(label_directory + r'/vali_location.npy', vali_location)
    # np.save(label_directory + r'/test_location.npy', test_location)
    # #保存每个wav的ID
    # np.save(label_directory + r'/train_id.npy',train_id)
    # np.save(label_directory + r'/vali_id.npy', vali_id)
    # np.save(label_directory + r'/test_id.npy', test_id)
    # #保存每个wav索引
    # train_index = get_index(train_data_directory)
    # vali_index = get_index(vali_data_directory)
    # test_index = get_index(test_data_directory)
    #
    # np.save(label_directory + r'/train_index.npy', train_index)
    # np.save(label_directory + r'/vali_index.npy', vali_index)
    # np.save(label_directory + r'/test_index.npy', test_index)



